import time
import serial
import serial.tools.list_ports
import threading
import os
import logging
from PyQt5.QtSerialPort import QSerialPort
from PyQt5.QtCore import QByteArray

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('fota_lib')

class FOTAUpgrader:
    """
    FOTA升级管理器类，用于管理FOTA升级过程
    """
    def __init__(self, com_port=None, baud_rate=115200, timeout=5, serial_port=None):
        """
        初始化FOTA升级管理器
        
        参数:
            com_port: 串口号，默认为None（自动检测）
            baud_rate: 波特率，默认为115200
            timeout: 超时时间，默认为5秒
            serial_port: 串口对象，如果提供则直接使用，不再重新打开串口
                         可以是PySerial的Serial对象或QSerialPort对象
        """
        self.com_port = com_port
        self.baud_rate = baud_rate
        self.timeout = timeout
        self.serial_port = serial_port  # 可以直接传入串口对象
        self.current_version = None
        self.new_version = None
        self.upgrade_status = None
        self.upgrade_progress = 0
        self.stop_event = threading.Event()
        self.external_serial = serial_port is not None  # 标记是否使用外部串口
        self.is_qserial = isinstance(serial_port, QSerialPort)  # 标记是否是QSerialPort
        
    # 修改write方法以支持QSerialPort
    def _write(self, data):
        """内部方法：写入数据到串口"""
        if self.is_qserial:
            if isinstance(data, str):
                data = data.encode('utf-8')
            if isinstance(data, bytes):
                qdata = QByteArray(data)
                return self.serial_port.write(qdata)
            return self.serial_port.write(data)
        else:
            return self.serial_port.write(data)
    
    # 修改read方法以支持QSerialPort
    def _read(self, size=1):
        """内部方法：从串口读取数据"""
        if self.is_qserial:
            if self.serial_port.waitForReadyRead(100):
                return self.serial_port.read(size).data()
            return b''
        else:
            return self.serial_port.read(size)
    
    # 修改readline方法以支持QSerialPort
    def _readline(self):
        """内部方法：从串口读取一行数据"""
        if self.is_qserial:
            buffer = b''
            while self.serial_port.waitForReadyRead(100):
                data = self.serial_port.readAll().data()
                buffer += data
                if b'\n' in data:
                    break
            return buffer
        else:
            return self.serial_port.readline()
    
    def open_serial_port(self):
        """
        打开串口
        
        返回:
            布尔值，表示是否成功打开串口
        """
        # 如果已经提供了串口对象，则直接返回成功
        if self.serial_port is not None:
            if self.is_qserial:
                if self.serial_port.isOpen():
                    logger.info(f"使用已提供的QSerialPort对象")
                    return True
            else:
                if self.serial_port.is_open:
                    logger.info(f"使用已提供的PySerial对象")
                    return True
        
        try:
            if self.com_port:
                self.serial_port = serial.Serial(
                    self.com_port, 
                    self.baud_rate, 
                    timeout=self.timeout
                )
                logger.info(f"成功打开串口 {self.com_port}，波特率: {self.baud_rate}")
                return True
            else:
                # 自动检测可用串口
                available_ports = self.get_available_ports()
                if available_ports:
                    self.com_port = available_ports[0]
                    self.serial_port = serial.Serial(
                        self.com_port, 
                        self.baud_rate, 
                        timeout=self.timeout
                    )
                    logger.info(f"成功打开串口 {self.com_port}，波特率: {self.baud_rate}")
                    return True
                else:
                    logger.error("未找到可用串口")
                    return False
        except serial.SerialException as e:
            logger.error(f"无法打开串口 {self.com_port}: {e}")
            return False
    
    def close_serial_port(self):
        """
        关闭串口
        """
        # 如果使用的是外部提供的串口，则不关闭
        if self.external_serial:
            logger.info("使用外部串口，不在此关闭")
            return
            
        if self.serial_port and self.serial_port.is_open:
            self.serial_port.close()
            logger.info(f"串口 {self.com_port} 已关闭")
    
    def get_available_ports(self):
        """
        获取所有可用的串口
        
        返回:
            可用串口列表
        """
        available_ports = []
        for port in serial.tools.list_ports.comports():
            try:
                ser = serial.Serial(port.device, self.baud_rate, timeout=self.timeout)
                ser.write(b'AT\r\n')
                response = ser.readline().decode().strip()
                if "OK" in response:
                    logger.info(f"端口 {port.device} 有响应，添加到可用列表")
                    available_ports.append(port.device)
                ser.close()
            except Exception as e:
                logger.error(f"检测端口 {port.device} 时出错: {e}")
        return available_ports
    
    def get_current_version(self):
        """
        获取当前固件版本
        
        返回:
            当前版本字符串，如果获取失败则返回None
        """
        if not self.serial_port:
            if not self.open_serial_port():
                return None
        
        # 检查串口是否打开，区分QSerialPort和PySerial
        if self.is_qserial:
            if not self.serial_port.isOpen():
                if not self.open_serial_port():
                    return None
        else:
            if not self.serial_port.is_open:
                if not self.open_serial_port():
                    return None
            
        try:
            self._write(b'AT+SGSW\r\n')
            time.sleep(1)  # 等待一段时间以获取返回
            response = self._readline()
            self.current_version = response.decode().strip() if isinstance(response, bytes) else response.strip()
            logger.info(f"当前版本: {self.current_version}")
            return self.current_version
        except Exception as e:
            logger.error(f"获取当前版本时出错: {e}")
            return None
    
    def start_upgrade(self, ftp_address):
        """
        开始FOTA升级
        
        参数:
            ftp_address: FTP地址，格式为 "ftp://用户名:密码@ftp地址/文件路径"
            
        返回:
            布尔值，表示是否成功启动升级
        """
        if not self.serial_port:
            if not self.open_serial_port():
                return False
        
        # 检查串口是否打开，区分QSerialPort和PySerial
        if self.is_qserial:
            if not self.serial_port.isOpen():
                if not self.open_serial_port():
                    return False
        else:
            if not self.serial_port.is_open:
                if not self.open_serial_port():
                    return False
            
        # 获取当前版本
        self.current_version = self.get_current_version()
        if not self.current_version:
            logger.error("无法获取当前版本，升级终止")
            return False
            
        # 发送升级命令
        try:
            self._write(f'AT+QFOTADL={ftp_address}\r\n'.encode())
            logger.info(f"已发送升级命令: AT+QFOTADL={ftp_address}")
            
            # 启动监控线程
            self.stop_event.clear()
            monitor_thread = threading.Thread(target=self._monitor_upgrade)
            monitor_thread.daemon = True
            monitor_thread.start()
            
            return True
        except Exception as e:
            logger.error(f"发送升级命令时出错: {e}")
            return False
    
    def _monitor_upgrade(self):
        """
        监控升级进度的内部方法
        """
        self.upgrade_status = "正在升级"
        self.upgrade_progress = 0
        
        while not self.stop_event.is_set():
            try:
                response = self._readline()
                if response:
                    logger.info(f"模块响应: {response}")
                    
                    if "FTPSTART" in response:
                        logger.info("升级开始...")
                        self.upgrade_status = "升级开始"
                        self.upgrade_progress = 10
                    elif "FTPEND,0" in response:
                        logger.info("升级成功...")
                        self.upgrade_status = "升级成功"
                        self.upgrade_progress = 90
                    elif "AT READY" in response:
                        logger.info("模组已开机，等待版本确认...")
                        self.upgrade_status = "等待版本确认"
                        self.upgrade_progress = 95
                        
                        # 发送AT+SGSW命令以获取新版本
                        self._write(b'AT+SGSW\r\n')
                        time.sleep(1)  # 等待一段时间以获取返回
                        self.new_version = self._readline().decode().strip()
                        logger.info(f"新版本: {self.new_version}")
                        
                        if self.new_version != self.current_version:
                            logger.info("版本确认，FOTA升级成功...")
                            self.upgrade_status = "升级完成"
                            self.upgrade_progress = 100
                            break
                    elif "ERROR" in response:
                        logger.error("升级过程中出错...")
                        self.upgrade_status = "升级失败"
                        break
                    elif "%" in response:
                        # 尝试解析进度百分比
                        try:
                            progress_parts = response.split("%")
                            if len(progress_parts) > 1:
                                progress_str = progress_parts[0].strip().split()[-1]
                                progress = int(progress_str)
                                self.upgrade_progress = 10 + int(progress * 0.8)  # 将进度映射到10%-90%
                                logger.info(f"升级进度: {self.upgrade_progress}%")
                        except Exception as e:
                            logger.error(f"解析进度时出错: {e}")
            except Exception as e:
                logger.error(f"监控升级进度时出错: {e}")
                time.sleep(0.5)
    
    def stop_upgrade(self):
        """
        停止升级监控
        """
        self.stop_event.set()
        logger.info("已停止升级监控")
    
    def get_upgrade_status(self):
        """
        获取升级状态
        
        返回:
            升级状态字符串
        """
        return self.upgrade_status
    
    def get_upgrade_progress(self):
        """
        获取升级进度
        
        返回:
            升级进度百分比
        """
        return self.upgrade_progress
    
    def is_upgrade_complete(self):
        """
        检查升级是否完成
        
        返回:
            布尔值，表示升级是否完成
        """
        return self.upgrade_status == "升级完成"
    
    def is_upgrade_failed(self):
        """
        检查升级是否失败
        
        返回:
            布尔值，表示升级是否失败
        """
        return self.upgrade_status == "升级失败"


class FOTATester:
    """
    FOTA测试器类，用于进行批量FOTA升级测试
    """
    def __init__(self, serial_ports=None):
        """
        初始化FOTA测试器
        
        参数:
            serial_ports: 串口对象列表，如果提供则直接使用，不再重新打开串口
        """
        self.serial_ports = serial_ports or []
        self.test_results = []
        self.external_serial = serial_ports is not None  # 标记是否使用外部串口
        
    def scan_ports(self, baud_rate=115200, timeout=5):
        """
        扫描可用串口
        
        参数:
            baud_rate: 波特率，默认为115200
            timeout: 超时时间，默认为5秒
            
        返回:
            可用串口列表
        """
        # 如果已经提供了串口对象，则直接返回
        if self.external_serial and self.serial_ports:
            return [ser.name for ser in self.serial_ports if hasattr(ser, 'name')]
            
        available_ports = serial.tools.list_ports.comports()
        self.serial_ports = []

        # 连接到每个串口并发送AT命令以检查响应
        for port in available_ports:
            try:
                ser = serial.Serial(port.device, baud_rate, timeout=timeout)
                ser.write(b'AT\r\n')
                response = ser.readline().decode().strip()
                if "OK" in response:
                    logger.info(f"端口 {port.device} 有响应，保留连接...")
                    self.serial_ports.append(ser)
                else:
                    logger.info(f"端口 {port.device} 无响应，删除连接...")
                    ser.close()
            except Exception as e:
                logger.error(f"连接端口 {port.device} 时出错: {e}")
                
        return [ser.name for ser in self.serial_ports]
    
    def run_test(self, ftp_address_A, ftp_address_B, cycles=10):
        """
        运行FOTA升级测试
        
        参数:
            ftp_address_A: A版本FTP地址
            ftp_address_B: B版本FTP地址
            cycles: 测试循环次数，默认为10
            
        返回:
            测试结果列表
        """
        if not self.serial_ports:
            logger.error("未找到可用串口，请先调用scan_ports()")
            return []
            
        self.test_results = []
        
        for i in range(cycles):
            logger.info(f"测试循环 {i+1} 开始...")
            for ser in self.serial_ports:
                logger.info(f"在端口 {ser.name} 上进行FOTA升级...")
                
                # 升级到B版本
                success_B = self._perform_upgrade(ser, ftp_address_B)
                
                # 升级到A版本
                success_A = self._perform_upgrade(ser, ftp_address_A)
                
                # 记录结果
                self.test_results.append((ser.name, i+1, success_B, success_A))
                
            logger.info(f"测试循环 {i+1} 结束...\n")
            
        # 输出测试结果
        logger.info("测试结果汇总:")
        for result in self.test_results:
            logger.info(f"端口: {result[0]}, 循环: {result[1]}, 升级B: {'成功' if result[2] else '失败'}, 升级A: {'成功' if result[3] else '失败'}")
            
        return self.test_results
    
    def _perform_upgrade(self, serial_port, ftp_address):
        """
        执行单次升级
        
        参数:
            serial_port: 串口对象
            ftp_address: FTP地址
            
        返回:
            布尔值，表示升级是否成功
        """
        try:
            # 发送AT+SGSW命令以获取当前版本
            serial_port.write(b'AT+SGSW\r\n')
            time.sleep(1)  # 等待一段时间以获取返回
            current_version = serial_port.readline().decode().strip()
            logger.info(f"当前版本: {current_version}")
            
            # 发送AT+QFOTADL=ftp地址命令以开始升级
            serial_port.write(f'AT+QFOTADL={ftp_address}\r\n'.encode())
            success = False
            
            # 设置超时时间
            timeout_start = time.time()
            timeout_limit = 300  # 5分钟超时
            
            while time.time() - timeout_start < timeout_limit:
                response = serial_port.readline().decode().strip()
                if response:
                    logger.info(f"模块响应: {response}")
                    
                    if "FTPSTART" in response:
                        logger.info("升级开始...")
                    elif "FTPEND,0" in response:
                        logger.info("升级成功...")
                        success = True
                        break
                    elif "AT READY" in response:
                        logger.info("模组已开机，等待版本确认...")
                        # 发送AT+SGSW命令以获取新版本
                        serial_port.write(b'AT+SGSW\r\n')
                        time.sleep(1)  # 等待一段时间以获取返回
                        new_version = serial_port.readline().decode().strip()
                        logger.info(f"新版本: {new_version}")
                        if new_version != current_version:
                            logger.info("版本确认，FOTA升级成功...")
                            success = True
                            break
                    elif "ERROR" in response:
                        logger.error("升级过程中出错...")
                        break
                        
                time.sleep(0.1)  # 短暂休眠，避免CPU占用过高
                
            if time.time() - timeout_start >= timeout_limit:
                logger.error("升级超时")
                
            return success
        except Exception as e:
            logger.error(f"执行升级时出错: {e}")
            return False
    
    def get_test_results(self):
        """
        获取测试结果
        
        返回:
            测试结果列表
        """
        return self.test_results
    
    def close_all_ports(self):
        """
        关闭所有串口
        """
        # 如果使用的是外部提供的串口，则不关闭
        if self.external_serial:
            logger.info("使用外部串口，不在此关闭")
            return
            
        for ser in self.serial_ports:
            try:
                if ser.is_open:
                    ser.close()
                    logger.info(f"串口 {ser.name} 已关闭")
            except Exception as e:
                logger.error(f"关闭串口 {ser.name} 时出错: {e}")
        
        self.serial_ports = []


def fota_upgrade(serial_port, ftp_address):
    """
    实现FOTA升级函数

    参数:
        serial_port: 串口实例，需要有write和readline方法
        ftp_address: FTP地址，格式为 "ftp://用户名:密码@ftp地址/文件路径"
        
    返回:
        升级是否成功
    """
    # 发送AT+SGSW命令以获取当前版本
    serial_port.write(b'AT+SGSW\r\n')
    time.sleep(1)  # 等待一段时间以获取返回
    current_version = serial_port.readline().decode().strip()
    logger.info(f"当前版本: {current_version}")
    
    # 发送AT+QFOTADL=ftp地址命令以开始升级
    serial_port.write(f'AT+QFOTADL={ftp_address}\r\n'.encode())
    success = False
    
    # 设置超时时间
    timeout_start = time.time()
    timeout_limit = 300  # 5分钟超时
    
    while time.time() - timeout_start < timeout_limit:
        response = serial_port.readline().decode().strip()
        if response:
            logger.info(f"模块响应: {response}")
            
            if "FTPSTART" in response:
                logger.info("升级开始...")
            elif "FTPEND,0" in response:
                logger.info("升级成功...")
                success = True
                break
            elif "AT READY" in response:
                logger.info("模组已开机，等待版本确认...")
                # 发送AT+SGSW命令以获取新版本
                serial_port.write(b'AT+SGSW\r\n')
                time.sleep(1)  # 等待一段时间以获取返回
                new_version = serial_port.readline().decode().strip()
                logger.info(f"新版本: {new_version}")
                if new_version != current_version:
                    logger.info("版本确认，FOTA升级成功...")
                    success = True
                    break
            elif "ERROR" in response:
                logger.error("升级过程中出错...")
                break
                
        time.sleep(0.1)  # 短暂休眠，避免CPU占用过高
        
    if time.time() - timeout_start >= timeout_limit:
        logger.error("升级超时")
        
    return success 